{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Transformer中的位置编码PE(position encoding)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_seq_len = 1000\n",
    "d_model = 128"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1.PositionalEncoding类函数编码"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PositionalEncoding(nn.Module):\n",
    "\n",
    "    def __init__(self, max_seq_len: int, d_model: int):\n",
    "        super().__init__()\n",
    "\n",
    "        # Assume d_model is an even number for convenience\n",
    "        assert d_model % 2 == 0   # 为了编码方便\n",
    "\n",
    "        # ---1.计算PE矩阵\n",
    "        # 位置编码二维矩阵PE的大小: [max_seq_len, d_model]\n",
    "        pe = torch.zeros(max_seq_len, d_model)  # 初始化为零矩阵\n",
    "        # 行：i向量 [0,1,2,..., 999]  表示每个时间步t\n",
    "        i_seq = torch.linspace(0, max_seq_len - 1, max_seq_len)\n",
    "        # 列：j向量 [0,2,4,6,8]       表示偶数位\n",
    "        j_seq = torch.linspace(0, d_model - 2, d_model // 2)   #(0, 8, 5)\n",
    "        # 生成网格数据: 2个矩阵[1000, 5]\n",
    "        pos, two_i = torch.meshgrid(i_seq, j_seq)\n",
    "        pe_2i = torch.sin(pos / 10000**(two_i / d_model))    # 偶数位sin\n",
    "        pe_2i_1 = torch.cos(pos / 10000**(two_i / d_model))  # 奇数位cos\n",
    "        # stack拼接到第2个维度[0,1,2]，在把3维重塑为2维\n",
    "        pe = torch.stack((pe_2i, pe_2i_1), 2).reshape(max_seq_len, d_model)\n",
    "\n",
    "        # ---2.定义嵌入层\n",
    "        self.embedding = nn.Embedding(max_seq_len, d_model)  # 定义了一个嵌入层\n",
    "        self.embedding.weight.data = pe      # 使用位置编码计算好的嵌入矩阵对其进行初始化\n",
    "        self.embedding.requires_grad_(False) # 将其参数设为不可训练\n",
    "\n",
    "    def forward(self, t):\n",
    "        # 调用嵌入层方法\n",
    "        # t表示抽取的时间点向量: [32, 43, 85, 31, 86, 90, 67, 61, 50, 33, 87, 48, 31, 48, 48, 93]\n",
    "        return self.embedding(t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/envs/d2l/lib/python3.9/site-packages/torch/functional.py:445: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  ../aten/src/ATen/native/TensorShape.cpp:2157.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "PositionalEncoding(\n",
       "  (embedding): Embedding(1000, 128)\n",
       ")"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "pe = PositionalEncoding(max_seq_len, d_model)\n",
    "pe"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "embedding.weight  : torch.Size([1000, 128])\n"
     ]
    }
   ],
   "source": [
    "# 方式一：查看模型参数返回\n",
    "for name, p in pe.named_parameters():\n",
    "    print(name, ' :', p.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "embedding.weight  : torch.Size([1000, 128])\n",
      "tensor([[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
      "          0.0000e+00,  1.0000e+00],\n",
      "        [ 8.4147e-01,  5.4030e-01,  7.6172e-01,  ...,  1.0000e+00,\n",
      "          1.1548e-04,  1.0000e+00],\n",
      "        [ 9.0930e-01, -4.1615e-01,  9.8705e-01,  ...,  1.0000e+00,\n",
      "          2.3096e-04,  1.0000e+00],\n",
      "        ...,\n",
      "        [-8.9797e-01, -4.4006e-01,  5.4092e-01,  ...,  9.9117e-01,\n",
      "          1.1488e-01,  9.9338e-01],\n",
      "        [-8.5547e-01,  5.1785e-01, -2.9021e-01,  ...,  9.9116e-01,\n",
      "          1.1499e-01,  9.9337e-01],\n",
      "        [-2.6461e-02,  9.9965e-01, -9.1697e-01,  ...,  9.9114e-01,\n",
      "          1.1511e-01,  9.9335e-01]])\n"
     ]
    }
   ],
   "source": [
    "# 方式二：查看模型参数返回\n",
    "for k, v in pe.state_dict().items():\n",
    "    print(k, ' :', v.shape)   # 参数名: 参数大小\n",
    "    print(v)  # 参数具体值"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### PE编码器的使用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 1])\n",
      "tensor([[40],\n",
      "        [ 0],\n",
      "        [ 6],\n",
      "        [52],\n",
      "        [47],\n",
      "        [71],\n",
      "        [96],\n",
      "        [99],\n",
      "        [52],\n",
      "        [ 6],\n",
      "        [98],\n",
      "        [47],\n",
      "        [57],\n",
      "        [17],\n",
      "        [ 5],\n",
      "        [52]])\n"
     ]
    }
   ],
   "source": [
    "n_steps = 100\n",
    "batch_size = 16\n",
    "t1 = torch.randint(0, n_steps, (batch_size, ))    # 随机抽取16个时间点\n",
    "t2 = torch.randint(0, n_steps, (batch_size, ))    # 随机抽取16个时间点\n",
    "t1 = t1.reshape(batch_size,1)\n",
    "t2 = t2.reshape(batch_size,1)\n",
    "print(t1.shape)\n",
    "print(t1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 1])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "t1.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 7.4511e-01, -6.6694e-01, -8.0965e-02,  ...,  9.9999e-01,\n",
       "           4.6191e-03,  9.9999e-01]],\n",
       "\n",
       "        [[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
       "           0.0000e+00,  1.0000e+00]],\n",
       "\n",
       "        [[-2.7942e-01,  9.6017e-01, -8.8542e-01,  ...,  1.0000e+00,\n",
       "           6.9287e-04,  1.0000e+00]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[-9.6140e-01, -2.7516e-01,  8.3414e-01,  ...,  1.0000e+00,\n",
       "           1.9631e-03,  1.0000e+00]],\n",
       "\n",
       "        [[-9.5892e-01,  2.8366e-01, -9.2771e-01,  ...,  1.0000e+00,\n",
       "           5.7739e-04,  1.0000e+00]],\n",
       "\n",
       "        [[ 9.8663e-01, -1.6299e-01,  8.6635e-01,  ...,  9.9998e-01,\n",
       "           6.0048e-03,  9.9998e-01]]])"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p1 = pe(t1)\n",
    "p1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 1, 128])"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p1.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### 与图片连接"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 先经过一个线性层，将词向量维度转换为与图片通道数一致"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "pe_dim = 128   # 词向量维度\n",
    "channel = 1    # 图片通道数"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "pe_linear = nn.Sequential(nn.Linear(pe_dim, channel), \n",
    "                           nn.ReLU(),\n",
    "                           nn.Linear(channel, channel))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 1]) torch.Size([16, 1, 128])\n"
     ]
    }
   ],
   "source": [
    "# 整体调用\n",
    "# 1. 先转换为pe\n",
    "p1 = pe(t1)   \n",
    "print(t1.shape, p1.shape)   # 1 -> 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 1, 1, 1])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 2. 经过线性层, 改变通道数\n",
    "n = 16   # 批量大小\n",
    "pe_v = pe_linear(p1).reshape(n, -1, 1, 1)  # (1, 128) -> (1,1) -> (1,1,1)\n",
    "# pe_linear: 降维 128 -> 1 通道数\n",
    "# reshape: 整理维度 [n, C, 1, 1]\n",
    "pe_v.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 逐行说明\n",
    "pl = pe_linear(p1)    \n",
    "# 输入编码后的向量：[16, 1, 128]    # 第2个参数代表图片的数量\n",
    "# 输出：仅改变通道数  [16, 1, 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 1, 1])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pl.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### 将整理后的位置编码与图片连接"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1, 28, 28])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randn(1, 28, 28) \n",
    "x.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 1, 1, 1])"
      ]
     },
     "execution_count": 121,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pe_v.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 1, 28, 28])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(x + pe_v).shape     # 相当于将原图的一个通道，变为了现在指定的16个通道\n",
    "\n",
    "# [16, 1, 28, 28]  \n",
    "# 16: 数据批量大小(时间点的个数 - batch_size)\n",
    "# 1: 通道数\n",
    "# 28*28: 图片大小"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2.逐行测试讲解"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.计算PE矩阵"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1000, 128])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 位置编码二维矩阵PE的大小: [max_seq_len, d_model]  1000 *128\n",
    "pe = torch.zeros(max_seq_len, d_model)  # 初始化为零矩阵\n",
    "pe.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([  0.,   1.,   2.,   3.,   4.,   5.,   6.,   7.,   8.,   9.,  10.,  11.,\n",
      "         12.,  13.,  14.,  15.,  16.,  17.,  18.,  19.,  20.,  21.,  22.,  23.,\n",
      "         24.,  25.,  26.,  27.,  28.,  29.,  30.,  31.,  32.,  33.,  34.,  35.,\n",
      "         36.,  37.,  38.,  39.,  40.,  41.,  42.,  43.,  44.,  45.,  46.,  47.,\n",
      "         48.,  49.,  50.,  51.,  52.,  53.,  54.,  55.,  56.,  57.,  58.,  59.,\n",
      "         60.,  61.,  62.,  63.,  64.,  65.,  66.,  67.,  68.,  69.,  70.,  71.,\n",
      "         72.,  73.,  74.,  75.,  76.,  77.,  78.,  79.,  80.,  81.,  82.,  83.,\n",
      "         84.,  85.,  86.,  87.,  88.,  89.,  90.,  91.,  92.,  93.,  94.,  95.,\n",
      "         96.,  97.,  98.,  99., 100., 101., 102., 103., 104., 105., 106., 107.,\n",
      "        108., 109., 110., 111., 112., 113., 114., 115., 116., 117., 118., 119.,\n",
      "        120., 121., 122., 123., 124., 125., 126., 127., 128., 129., 130., 131.,\n",
      "        132., 133., 134., 135., 136., 137., 138., 139., 140., 141., 142., 143.,\n",
      "        144., 145., 146., 147., 148., 149., 150., 151., 152., 153., 154., 155.,\n",
      "        156., 157., 158., 159., 160., 161., 162., 163., 164., 165., 166., 167.,\n",
      "        168., 169., 170., 171., 172., 173., 174., 175., 176., 177., 178., 179.,\n",
      "        180., 181., 182., 183., 184., 185., 186., 187., 188., 189., 190., 191.,\n",
      "        192., 193., 194., 195., 196., 197., 198., 199., 200., 201., 202., 203.,\n",
      "        204., 205., 206., 207., 208., 209., 210., 211., 212., 213., 214., 215.,\n",
      "        216., 217., 218., 219., 220., 221., 222., 223., 224., 225., 226., 227.,\n",
      "        228., 229., 230., 231., 232., 233., 234., 235., 236., 237., 238., 239.,\n",
      "        240., 241., 242., 243., 244., 245., 246., 247., 248., 249., 250., 251.,\n",
      "        252., 253., 254., 255., 256., 257., 258., 259., 260., 261., 262., 263.,\n",
      "        264., 265., 266., 267., 268., 269., 270., 271., 272., 273., 274., 275.,\n",
      "        276., 277., 278., 279., 280., 281., 282., 283., 284., 285., 286., 287.,\n",
      "        288., 289., 290., 291., 292., 293., 294., 295., 296., 297., 298., 299.,\n",
      "        300., 301., 302., 303., 304., 305., 306., 307., 308., 309., 310., 311.,\n",
      "        312., 313., 314., 315., 316., 317., 318., 319., 320., 321., 322., 323.,\n",
      "        324., 325., 326., 327., 328., 329., 330., 331., 332., 333., 334., 335.,\n",
      "        336., 337., 338., 339., 340., 341., 342., 343., 344., 345., 346., 347.,\n",
      "        348., 349., 350., 351., 352., 353., 354., 355., 356., 357., 358., 359.,\n",
      "        360., 361., 362., 363., 364., 365., 366., 367., 368., 369., 370., 371.,\n",
      "        372., 373., 374., 375., 376., 377., 378., 379., 380., 381., 382., 383.,\n",
      "        384., 385., 386., 387., 388., 389., 390., 391., 392., 393., 394., 395.,\n",
      "        396., 397., 398., 399., 400., 401., 402., 403., 404., 405., 406., 407.,\n",
      "        408., 409., 410., 411., 412., 413., 414., 415., 416., 417., 418., 419.,\n",
      "        420., 421., 422., 423., 424., 425., 426., 427., 428., 429., 430., 431.,\n",
      "        432., 433., 434., 435., 436., 437., 438., 439., 440., 441., 442., 443.,\n",
      "        444., 445., 446., 447., 448., 449., 450., 451., 452., 453., 454., 455.,\n",
      "        456., 457., 458., 459., 460., 461., 462., 463., 464., 465., 466., 467.,\n",
      "        468., 469., 470., 471., 472., 473., 474., 475., 476., 477., 478., 479.,\n",
      "        480., 481., 482., 483., 484., 485., 486., 487., 488., 489., 490., 491.,\n",
      "        492., 493., 494., 495., 496., 497., 498., 499., 500., 501., 502., 503.,\n",
      "        504., 505., 506., 507., 508., 509., 510., 511., 512., 513., 514., 515.,\n",
      "        516., 517., 518., 519., 520., 521., 522., 523., 524., 525., 526., 527.,\n",
      "        528., 529., 530., 531., 532., 533., 534., 535., 536., 537., 538., 539.,\n",
      "        540., 541., 542., 543., 544., 545., 546., 547., 548., 549., 550., 551.,\n",
      "        552., 553., 554., 555., 556., 557., 558., 559., 560., 561., 562., 563.,\n",
      "        564., 565., 566., 567., 568., 569., 570., 571., 572., 573., 574., 575.,\n",
      "        576., 577., 578., 579., 580., 581., 582., 583., 584., 585., 586., 587.,\n",
      "        588., 589., 590., 591., 592., 593., 594., 595., 596., 597., 598., 599.,\n",
      "        600., 601., 602., 603., 604., 605., 606., 607., 608., 609., 610., 611.,\n",
      "        612., 613., 614., 615., 616., 617., 618., 619., 620., 621., 622., 623.,\n",
      "        624., 625., 626., 627., 628., 629., 630., 631., 632., 633., 634., 635.,\n",
      "        636., 637., 638., 639., 640., 641., 642., 643., 644., 645., 646., 647.,\n",
      "        648., 649., 650., 651., 652., 653., 654., 655., 656., 657., 658., 659.,\n",
      "        660., 661., 662., 663., 664., 665., 666., 667., 668., 669., 670., 671.,\n",
      "        672., 673., 674., 675., 676., 677., 678., 679., 680., 681., 682., 683.,\n",
      "        684., 685., 686., 687., 688., 689., 690., 691., 692., 693., 694., 695.,\n",
      "        696., 697., 698., 699., 700., 701., 702., 703., 704., 705., 706., 707.,\n",
      "        708., 709., 710., 711., 712., 713., 714., 715., 716., 717., 718., 719.,\n",
      "        720., 721., 722., 723., 724., 725., 726., 727., 728., 729., 730., 731.,\n",
      "        732., 733., 734., 735., 736., 737., 738., 739., 740., 741., 742., 743.,\n",
      "        744., 745., 746., 747., 748., 749., 750., 751., 752., 753., 754., 755.,\n",
      "        756., 757., 758., 759., 760., 761., 762., 763., 764., 765., 766., 767.,\n",
      "        768., 769., 770., 771., 772., 773., 774., 775., 776., 777., 778., 779.,\n",
      "        780., 781., 782., 783., 784., 785., 786., 787., 788., 789., 790., 791.,\n",
      "        792., 793., 794., 795., 796., 797., 798., 799., 800., 801., 802., 803.,\n",
      "        804., 805., 806., 807., 808., 809., 810., 811., 812., 813., 814., 815.,\n",
      "        816., 817., 818., 819., 820., 821., 822., 823., 824., 825., 826., 827.,\n",
      "        828., 829., 830., 831., 832., 833., 834., 835., 836., 837., 838., 839.,\n",
      "        840., 841., 842., 843., 844., 845., 846., 847., 848., 849., 850., 851.,\n",
      "        852., 853., 854., 855., 856., 857., 858., 859., 860., 861., 862., 863.,\n",
      "        864., 865., 866., 867., 868., 869., 870., 871., 872., 873., 874., 875.,\n",
      "        876., 877., 878., 879., 880., 881., 882., 883., 884., 885., 886., 887.,\n",
      "        888., 889., 890., 891., 892., 893., 894., 895., 896., 897., 898., 899.,\n",
      "        900., 901., 902., 903., 904., 905., 906., 907., 908., 909., 910., 911.,\n",
      "        912., 913., 914., 915., 916., 917., 918., 919., 920., 921., 922., 923.,\n",
      "        924., 925., 926., 927., 928., 929., 930., 931., 932., 933., 934., 935.,\n",
      "        936., 937., 938., 939., 940., 941., 942., 943., 944., 945., 946., 947.,\n",
      "        948., 949., 950., 951., 952., 953., 954., 955., 956., 957., 958., 959.,\n",
      "        960., 961., 962., 963., 964., 965., 966., 967., 968., 969., 970., 971.,\n",
      "        972., 973., 974., 975., 976., 977., 978., 979., 980., 981., 982., 983.,\n",
      "        984., 985., 986., 987., 988., 989., 990., 991., 992., 993., 994., 995.,\n",
      "        996., 997., 998., 999.]) tensor([  0.,   2.,   4.,   6.,   8.,  10.,  12.,  14.,  16.,  18.,  20.,  22.,\n",
      "         24.,  26.,  28.,  30.,  32.,  34.,  36.,  38.,  40.,  42.,  44.,  46.,\n",
      "         48.,  50.,  52.,  54.,  56.,  58.,  60.,  62.,  64.,  66.,  68.,  70.,\n",
      "         72.,  74.,  76.,  78.,  80.,  82.,  84.,  86.,  88.,  90.,  92.,  94.,\n",
      "         96.,  98., 100., 102., 104., 106., 108., 110., 112., 114., 116., 118.,\n",
      "        120., 122., 124., 126.])\n"
     ]
    }
   ],
   "source": [
    "# 行：i向量 [0,1,2,..., 999]       表示每个时间步t\n",
    "i_seq = torch.linspace(0, max_seq_len - 1, max_seq_len)\n",
    "# 列：j向量 [0,2,4,6,8, 126]       表示位置编码向量的偶数位\n",
    "j_seq = torch.linspace(0, d_model - 2, d_model // 2)    #64\n",
    "print(i_seq,j_seq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "64"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(j_seq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 生成网格数据: 2个矩阵[1000, 64]\n",
    "pos, two_i = torch.meshgrid(i_seq, j_seq)\n",
    "pe_2i = torch.sin(pos / 10000**(two_i / d_model))    # 偶数位sin\n",
    "pe_2i_1 = torch.cos(pos / 10000**(two_i / d_model))  # 奇数位cos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1000, 64])"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pe_2i.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1000, 64, 2])"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.stack((pe_2i, pe_2i_1), 2).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# stack拼接到第2个维度[0,1,2]，在把3维重塑为2维\n",
    "pe = torch.stack((pe_2i, pe_2i_1), 2).reshape(max_seq_len, d_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([1000, 128])"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pe.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 0.0000e+00,  1.0000e+00,  0.0000e+00,  ...,  1.0000e+00,\n",
       "          0.0000e+00,  1.0000e+00],\n",
       "        [ 8.4147e-01,  5.4030e-01,  7.6172e-01,  ...,  1.0000e+00,\n",
       "          1.1548e-04,  1.0000e+00],\n",
       "        [ 9.0930e-01, -4.1615e-01,  9.8705e-01,  ...,  1.0000e+00,\n",
       "          2.3096e-04,  1.0000e+00],\n",
       "        ...,\n",
       "        [-8.9797e-01, -4.4006e-01,  5.4092e-01,  ...,  9.9117e-01,\n",
       "          1.1488e-01,  9.9338e-01],\n",
       "        [-8.5547e-01,  5.1785e-01, -2.9021e-01,  ...,  9.9116e-01,\n",
       "          1.1499e-01,  9.9337e-01],\n",
       "        [-2.6461e-02,  9.9965e-01, -9.1697e-01,  ...,  9.9114e-01,\n",
       "          1.1511e-01,  9.9335e-01]])"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pe"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.Embedding结构\n",
    "\n",
    "- Embedding.weight会从*标准正态分布*中初始化成大小为（num_embeddings, embedding_dim）的矩阵\n",
    "    - PE矩阵的作用就是替换这个标准正态分布\n",
    "- input中的标号表示从矩阵对应行获取权重来表示单词\n",
    "\n",
    "参考：http://t.csdnimg.cn/ZCKQd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_seq_len = 1000   # 句子的最大长度\n",
    "d_model = 128        # 词向量的长度"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 设置embedding结构\n",
    "e = nn.Embedding(max_seq_len, d_model) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 输入的时间向量\n",
    "t2 = torch.randint(0, n_steps, (batch_size, ))   # 一句话16个词\n",
    "t3 = torch.randint(0, n_steps, (batch_size, 2))  # 2句话"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16]) torch.Size([16, 2])\n"
     ]
    }
   ],
   "source": [
    "print(t2.shape, t3.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 128])"
      ]
     },
     "execution_count": 61,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# embedding后：每行就是一个词向量\n",
    "e2 = e(t2)\n",
    "e2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([16, 2, 128])"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "e3 = e(t3)\n",
    "e3.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.8290,  1.3628,  0.5724,  ...,  0.3281, -1.2781,  0.8026],\n",
       "        [-0.8468,  0.3652,  0.1609,  ...,  1.0198,  1.3574,  0.3451],\n",
       "        [-0.3440, -0.4354,  0.3475,  ...,  1.3262, -0.1968,  0.4624],\n",
       "        ...,\n",
       "        [-0.9771,  1.6782,  2.0468,  ...,  0.5930,  1.6419, -0.0021],\n",
       "        [ 0.3120, -1.4649,  1.7905,  ...,  1.2019,  1.5737,  1.2519],\n",
       "        [-0.2558, -0.2975, -1.0529,  ..., -1.8280, -0.9493, -0.9284]])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# embedding的权重大小\n",
    "e.weight.data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[True, True, True,  ..., True, True, True],\n",
       "         [True, True, True,  ..., True, True, True]],\n",
       "\n",
       "        [[True, True, True,  ..., True, True, True],\n",
       "         [True, True, True,  ..., True, True, True]],\n",
       "\n",
       "        [[True, True, True,  ..., True, True, True],\n",
       "         [True, True, True,  ..., True, True, True]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[True, True, True,  ..., True, True, True],\n",
       "         [True, True, True,  ..., True, True, True]],\n",
       "\n",
       "        [[True, True, True,  ..., True, True, True],\n",
       "         [True, True, True,  ..., True, True, True]],\n",
       "\n",
       "        [[True, True, True,  ..., True, True, True],\n",
       "         [True, True, True,  ..., True, True, True]]])"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 从权重中取出对应索引的数值\n",
    "e.weight.data[t3] == e3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.8293,  0.8427, -0.5874,  ...,  0.7140,  0.2557, -0.3822],\n",
       "         [ 1.1392, -0.1953, -0.7295,  ...,  0.8760, -0.4759,  2.1021]],\n",
       "\n",
       "        [[ 1.6802, -0.3143, -0.8238,  ..., -1.2749, -0.0555,  0.4325],\n",
       "         [ 1.6263,  0.1748, -0.0375,  ..., -2.0088, -0.2860,  1.4281]],\n",
       "\n",
       "        [[-0.0757, -0.2243,  1.3649,  ..., -0.3594, -0.9480, -1.0630],\n",
       "         [-0.8829,  0.1561,  0.2510,  ..., -0.6215, -1.1676, -0.7077]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[ 0.0803,  1.1545, -0.7426,  ...,  0.1323,  0.5377,  0.2642],\n",
       "         [ 0.8365,  1.1867, -0.1665,  ...,  0.8210,  0.0919,  1.8184]],\n",
       "\n",
       "        [[-0.3055,  0.0439, -1.3504,  ...,  1.7765,  0.2634, -0.5820],\n",
       "         [ 1.1935,  0.5321, -1.0562,  ...,  0.0259,  0.2304,  0.9908]],\n",
       "\n",
       "        [[ 0.3471, -0.3673,  0.1638,  ...,  0.2270, -0.0317,  0.4311],\n",
       "         [-0.8829,  0.1561,  0.2510,  ..., -0.6215, -1.1676, -0.7077]]],\n",
       "       grad_fn=<EmbeddingBackward0>)"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "e3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.8293,  0.8427, -0.5874,  ...,  0.7140,  0.2557, -0.3822],\n",
       "         [ 1.1392, -0.1953, -0.7295,  ...,  0.8760, -0.4759,  2.1021]],\n",
       "\n",
       "        [[ 1.6802, -0.3143, -0.8238,  ..., -1.2749, -0.0555,  0.4325],\n",
       "         [ 1.6263,  0.1748, -0.0375,  ..., -2.0088, -0.2860,  1.4281]],\n",
       "\n",
       "        [[-0.0757, -0.2243,  1.3649,  ..., -0.3594, -0.9480, -1.0630],\n",
       "         [-0.8829,  0.1561,  0.2510,  ..., -0.6215, -1.1676, -0.7077]],\n",
       "\n",
       "        ...,\n",
       "\n",
       "        [[ 0.0803,  1.1545, -0.7426,  ...,  0.1323,  0.5377,  0.2642],\n",
       "         [ 0.8365,  1.1867, -0.1665,  ...,  0.8210,  0.0919,  1.8184]],\n",
       "\n",
       "        [[-0.3055,  0.0439, -1.3504,  ...,  1.7765,  0.2634, -0.5820],\n",
       "         [ 1.1935,  0.5321, -1.0562,  ...,  0.0259,  0.2304,  0.9908]],\n",
       "\n",
       "        [[ 0.3471, -0.3673,  0.1638,  ...,  0.2270, -0.0317,  0.4311],\n",
       "         [-0.8829,  0.1561,  0.2510,  ..., -0.6215, -1.1676, -0.7077]]])"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "e.weight.data[t3] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "d2l",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
